from functools import wraps
from datetime import datetime, timedelta
import jwt
from flask import request, jsonify
from demoPython.config import JWT_SECRET_KEY

def generate_token(user_id):
    """生成JWT token"""
    payload = {
        'user_id': user_id,
        'exp': datetime.utcnow() + timedelta(days=1)  # token有效期1天
    }
    return jwt.encode(payload, JWT_SECRET_KEY, algorithm='HS256')

def verify_token(token):
    """验证JWT token"""
    try:
        payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=['HS256'])
        return payload['user_id']
    except jwt.ExpiredSignatureError:
        return None
    except jwt.InvalidTokenError:
        return None

def login_required(f):
    """登录验证装饰器"""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        token = None
        # 从请求头中获取token
        auth_header = request.headers.get('Authorization')
        if auth_header and auth_header.startswith('Bearer '):
            token = auth_header.split(' ')[1]
        
        # 从cookie中获取token
        if not token:
            token = request.cookies.get('token')
            
        if not token:
            return jsonify({'message': '未登录', 'code': 401}), 401
            
        user_id = verify_token(token)
        if not user_id:
            return jsonify({'message': 'token已过期或无效', 'code': 401}), 401
            
        # 将user_id添加到请求上下文中
        request.user_id = user_id
        return f(*args, **kwargs)
    return decorated_function